'''
Created on Sep 9, 2009

@author: mkiyer
'''

from optparse import OptionParser
import sys
import os
import logging
import h5py
import numpy as np
import collections
import re
import string

from veggie.genome.chrom import get_chrom_names

class Gene(object):    
    def __init__(self):
        self.exons = []
        self.introns = []
        self.utr5 = None
        self.utr3 = None
        self.aliases = []
        self.ncbi_id = None

def parse_kgxref(genemap, fhd):
    genelookup = collections.defaultdict(lambda: [])
    for line in fhd:
        fields = line.strip().split('\t')        
        if fields[0] in genemap:
            g = genemap[fields[0]]
            genelookup[fields[0]] = g            
            genelookup[fields[4]].append(g)
            g.aliases.append(fields[4].replace('/', '_'))
            #for field in fields[1:]:
            #    genelookup[field].append(g)
            #    g.aliases.append(field)
    return genelookup


def parse_reflink(acc_to_gene, fhd):    
    for line in fhd:
        if line.startswith('#'):
            continue
        fields = line.strip().split('\t')
        acc = fields[2]
        if acc in acc_to_gene:
            for g in acc_to_gene[acc]:
                g.ncbi_id = int(fields[6])


def parse_refseq_genes(fhd):
    '''
    hugo name varchar(255) NOT NULL default '',
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        fields = line.strip().split('\t')
        g = Gene()
        g.symbol = fields[0]
        g.acc = fields[1]
        g.chrom = fields[2]
        if g.chrom not in get_chrom_names():
            continue
        g.strand = fields[3]
        g.txstart = int(fields[4])
        g.txend = int(fields[5])
        g.cdsstart = int(fields[6])
        g.cdsend = int(fields[7])
        g.exon_count = int(fields[8])
        g.exon_starts = map(int, fields[9].split(',')[:-1])
        g.exon_ends = map(int, fields[10].split(',')[:-1])
        g.exons = zip(g.exon_starts, g.exon_ends)
        g.introns = zip(g.exon_ends, g.exon_starts[1:])
        
        utr5_exons = []
        utr3_exons = []
        for e in g.exons:
            if e[0] < g.cdsstart:
                utr5_exons.append(e)
            if e[1] > g.cdsend:
                utr3_exons.append(e)
        if g.strand == '+':
            g.utr5_exons = utr5_exons
            g.utr3_exons = utr3_exons
        else:
            g.utr5_exons = utr3_exons
            g.utr3_exons = utr5_exons

        # figure out the UTRs
        if g.txstart != g.cdsstart:
            if g.strand == '+':
                g.utr5 = (g.txstart, g.cdsstart)
            elif g.strand == '-':
                g.utr3 = (g.txstart, g.cdsstart)
            else:
                logging.critical("Gene strand is not '+' or '-': %s" % g.name)
                sys.exit(1)
        if g.txend != g.cdsend:
            if g.strand == '+':
                g.utr3 = (g.cdsend, g.txend)
            elif g.strand == '-':
                g.utr5 = (g.cdsend, g.txend)
            else:
                logging.critical("Gene strand is not '+' or '-': %s" % g.name)
                sys.exit(1)
        yield g

def parse_ucsc_genes(fhd):
    '''
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        fields = line.strip().split('\t')
        g = Gene()
        g.name = fields[0]
        g.chrom = fields[1]
        if g.chrom not in hg18_chrom_lengths:
            continue
        g.strand = fields[2]
        g.txstart = int(fields[3])
        g.txend = int(fields[4])
        g.cdsstart = int(fields[5])
        g.cdsend = int(fields[6])
        g.exon_count = int(fields[7])
        g.exon_starts = map(int, fields[8].split(',')[:-1])
        g.exon_ends = map(int, fields[9].split(',')[:-1])
        g.exons = zip(g.exon_starts, g.exon_ends)
        g.introns = zip(g.exon_ends, g.exon_starts[1:])
        # figure out the UTRs
        if g.txstart != g.cdsstart:
            if g.strand == '+':
                g.utr5 = (g.txstart, g.cdsstart)
            elif g.strand == '-':
                g.utr3 = (g.txstart, g.cdsstart)
            else:
                logging.critical("Gene strand is not '+' or '-': %s" % g.name)
                sys.exit(1)
        if g.txend != g.cdsend:
            if g.strand == '+':
                g.utr3 = (g.cdsend, g.txend)
            elif g.strand == '-':
                g.utr5 = (g.cdsend, g.txend)
            else:
                logging.critical("Gene strand is not '+' or '-': %s" % g.name)
                sys.exit(1)
        yield g

def get_exon_to_gene_map(genes):
    exonmap = collections.defaultdict(lambda: [])
    for g in genes:
        for start, end in g.exons:
            exonmap[(g.chrom, start, end)].append(g)
    return exonmap

def print_result_header():
    sys.stdout.write('Gene Name\tAccession\tNCBI Gene ID\tChrom\tStart\tEnd\tStrand\tGenomic Location\n')

def print_result_line(g, start, end, label):
    sys.stdout.write('%s\t%s\t%s\t%s\t%d\t%d\t%s\t%s\t%s\n' %
                     (g.symbol,
                      g.acc,
                      g.ncbi_id,
                      g.chrom,
                      start,
                      end,
                      g.strand,
                      '%s:%d-%d' % (g.chrom, start, end),
                      label))
                    
def get_full_genes(sym_to_gene, fhd):
    for line in fhd:
        gene_sym = line.strip()
        genes = sym_to_gene.get(gene_sym, None)
        if genes == None:
            print '%s\tError, could not find in RefSeq gene list' % (gene_sym)
            continue
        exons = set([])
        for g in genes:
            for ei, e in enumerate(g.exons):                
                if (g.chrom, e[0], e[1]) not in exons:
                    print_result_line(g, e[0], e[1],
                                      'exon_%d_of_%d' % (ei, len(g.exons)))
                    exons.add((g.chrom, e[0], e[1]))

def get_utrs(sym_to_gene, fhd):
    for line in fhd:
        gene_sym = line.strip()
        genes = sym_to_gene.get(gene_sym, None)
        if genes == None:
            print '%s\tError, could not find in RefSeq gene list' % (gene_sym)
            continue
        utrs = set([])
        for g in genes:
            for e in g.utr3_exons:
                if (g.chrom, e[0], e[1]) not in utrs:
                    print_result_line(g, e[0], e[1], '3UTR')
                    utrs.add((g.chrom, e[0], e[1]))
            for e in g.utr5_exons:
                if (g.chrom, e[0], e[1]) not in utrs:
                    print_result_line(g, e[0], e[1], '5UTR')
                    utrs.add((g.chrom, e[0], e[1]))

if __name__ == '__main__':    
    input_file = sys.argv[1]
    # read refseq genes    
    logging.basicConfig(level=logging.DEBUG)
    logging.debug("loading refseq genes")
    genes = [x for x in parse_refseq_genes(open('refFlat.txt', 'rU'))]    
    sym_to_gene = collections.defaultdict(lambda: [])
    acc_to_gene = collections.defaultdict(lambda: [])
    for g in genes:
        sym_to_gene[g.symbol].append(g)
        acc_to_gene[g.acc].append(g)
    logging.debug('parsing ncbi gene ids')
    parse_reflink(acc_to_gene, open('refLink.txt'))

    # print header
    print_result_header()
    # process input genes
    get_full_genes(sym_to_gene, open(input_file, 'rU'))
    #get_utrs(sym_to_gene, open(input_file, 'rU'))



